﻿using System;
using System.Collections.Generic;
using System.Text;
using Common.Utilities;

namespace Workflows.Components.TextMining.Clustering
{
    /// <summary>
    /// 
    /// </summary>
    public class KMeansCluster:IDocumentCluster
    {
        #region fields
        private Dictionary<int, Dictionary<int, DocTermFreq>> _DocTerms;
        private Dictionary<string, List<int>> _TermStemmings2;
        private Dictionary<int, string> _TermStemmings1;
        private int _NGramLen = 0;
        private IDistanceCalculator _DistanceCalculator;
        private IPurityCalculator _PurityCalculator;
        private Dictionary<int, double> _TermWeights;
        private Dictionary<int, ClusterNode> _ClusterResult;

        private double _ReassociationDistanceRatio = 1.50;
        private bool _CanOverlap;
        private int _K = 8;
        private const int _MaxIterationCount = 100;
        private double _TermFreqThreshold;
        private double _TermCoverageThreshold;
        private double _ClusterDistanceThreshold;
        private Dictionary<int, double> _DocDistances;              // docid -> distance (average of distances to other docs within current cluster)
        #endregion

        #region props
        
        /// <summary></summary>
        public int K
        {
            get
            {
                return this._K;
            }
            set
            {
                this._K = value;
            }
        }
        #endregion

        #region IDocumentCluster Members

        public event ClusterStatusChanged OnClusterStatusChanged;

        public IDistanceCalculator DistanceCalculator
        {
            get { return _DistanceCalculator; }
            set { _DistanceCalculator=value; }
        }

        public IPurityCalculator PurityCalculator
        {
            get { return _PurityCalculator; }
            set { _PurityCalculator=value; }
        }

        public Dictionary<int, double> TermWeights
        {
            get { return _TermWeights; }
            set { _TermWeights=value; }
        }

        public bool CanOverlap
        {
            get { return _CanOverlap; }
            set { _CanOverlap = value; }
        }

        public double ReassociationRange
        {
            get { return _ReassociationDistanceRatio; }
            set { this._ReassociationDistanceRatio = value; }
        }

        public void SetThreshold(double termCoverageThreshold, double termFreqThreshold, double clusterDistanceThreshold)
        {
            this._TermCoverageThreshold = termCoverageThreshold;
            this._TermFreqThreshold = termFreqThreshold;
            this._ClusterDistanceThreshold = clusterDistanceThreshold;
        }

        public Dictionary<int, ClusterNode> ClusterDocuments(
            Dictionary<int, Dictionary<int, DocTermFreq>> docTerms,
            bool useNgram, int ngramLen)
        {
            if(useNgram)
                this._NGramLen = ngramLen;
            return this.ClusterDocuments(docTerms, null, null);
        }

        public Dictionary<int, ClusterNode> ClusterDocuments(
            Dictionary<int, Dictionary<int, DocTermFreq>> docTerms,
            Dictionary<int,string> termStemmings1,
            Dictionary<string,List<int>> termStemmings2)
        {
            this._DocTerms = docTerms;
            this._TermStemmings1 = termStemmings1;
            this._TermStemmings2 = termStemmings2;

            if (this._DistanceCalculator == null)
                throw new Exception("Calculator must be initialized first before calculation");

            this._ClusterResult = new Dictionary<int, ClusterNode>();
            this.InitKClusters();
            int iterationCount = 0;
            int docsReassociated = this.UpdateDocAssociations();
            while (docsReassociated > 0)
            {
                foreach(int clusterID in this._ClusterResult.Keys)
                {
                    ClusterNode clusterNode = this._ClusterResult[clusterID];
                    this.CalculateClusterDistance(ref clusterNode);
                }

                if(this.OnClusterStatusChanged!=null)
                {
                    string msg = string.Format("K-means clustering: {0} docs reassociated...", docsReassociated);
                    int pct = iterationCount*100/_MaxIterationCount;
                    this.OnClusterStatusChanged(this._ClusterResult.Count, 0, msg, pct);
                }
                iterationCount += 1;

                if(iterationCount>=_MaxIterationCount)
                    break;

                docsReassociated = this.UpdateDocAssociations();
            }

            if (this.CanOverlap)
            {
                iterationCount = 0;
                docsReassociated = this.ReassociateDocs();
                while (docsReassociated > 0)
                {
                    foreach (int clusterID in this._ClusterResult.Keys)
                    {
                        ClusterNode clusterNode = this._ClusterResult[clusterID];
                        this.CalculateClusterDistance(ref clusterNode);
                    }

                    if (this.OnClusterStatusChanged != null)
                    {
                        string msg = string.Format("K-means clustering: {0} docs reassociated...", docsReassociated);
                        int pct = iterationCount*100/_MaxIterationCount;
                        this.OnClusterStatusChanged(this._ClusterResult.Count, 0, msg, pct);
                    }
                    iterationCount += 1;

                    if (iterationCount >= _MaxIterationCount)
                        break;

                    docsReassociated = this.ReassociateDocs();
                }
            }

            return this._ClusterResult;
        }

        #endregion

        #region init
        private void InitKClusters()
        {
            int[] docIDs=new int[this._DocTerms.Count];
            this._DocTerms.Keys.CopyTo(docIDs,0);
            this._DocDistances=new Dictionary<int, double>();
            this._ClusterResult=new Dictionary<int, ClusterNode>();
            int initialItemCountPerCluster = Math.Max(1, docIDs.Length / this._K);
            int i = 0;
            ClusterNode clusterNode=new ClusterNode();
            clusterNode.ClusterID = this.NextClusterID();
            clusterNode.Distance = 1;
            clusterNode.DocIDs=new List<int>();

            while (i<docIDs.Length)
            {
                clusterNode.DocIDs.Add(docIDs[i]);

                if(clusterNode.DocIDs.Count >= initialItemCountPerCluster)
                {
                    this.CalculateClusterDistance(ref clusterNode);
                    this._ClusterResult.Add(clusterNode.ClusterID,clusterNode);
                    clusterNode = new ClusterNode();
                    clusterNode.ClusterID = this.NextClusterID();
                    clusterNode.Distance = 1;
                    clusterNode.DocIDs = new List<int>();
                }
                i++;
            }
            if(clusterNode.DocIDs.Count>0)
            {
                this.CalculateClusterDistance(ref clusterNode);
                this._ClusterResult.Add(clusterNode.ClusterID,clusterNode);
            }
        }

        private void CalculateClusterDistance(ref ClusterNode clusterNode)
        {
            List<double> clusterDistances=new List<double>();
            Dictionary<int,double> supportiveTermWeights=new Dictionary<int, double>();
            for(int i=0;i<clusterNode.DocIDs.Count;i++)
            {
                int docID1 = clusterNode.DocIDs[i];
                List<double> docDistances=new List<double>();
                Dictionary<int, DocTermFreq> termFreqs1 = this._DocTerms[clusterNode.DocIDs[i]];
                //Dictionary<int, int> termFreq1 = this.InitTermFrequency(termFreqs1);
                for(int k=0;k<clusterNode.DocIDs.Count;k++)
                {
                    if (k != i)
                    {
                        // int docID2 = clusterNode.DocIDs[k];
                        Dictionary<int, DocTermFreq> termFreqs2 = this._DocTerms[clusterNode.DocIDs[k]];
                        //Dictionary<int, int> termFreq2 = this.InitTermFrequency(termFreqs2);
                        Dictionary<int, int> sharedTermIDs = new Dictionary<int, int>();
                        double distance=0;
                        if (this._TermWeights == null || this._TermWeights.Count == 0)
                            distance =
                                this._DistanceCalculator.CalculateDistance(
                                    termFreqs1, termFreqs2, ref sharedTermIDs);
                        else if (this._TermStemmings1 == null || this._TermStemmings2 == null)
                            distance = this._DistanceCalculator.CalculateDistance(
                                termFreqs1, termFreqs2, this._TermWeights,
                                true, this._NGramLen, ref sharedTermIDs);
                        else
                            distance = this._DistanceCalculator.CalculateDistance(
                                termFreqs1, termFreqs2, TermWeights,
                                this._TermStemmings1, this._TermStemmings2, ref sharedTermIDs);
                        clusterDistances.Add(distance);
                        docDistances.Add(distance);

                        if(sharedTermIDs.Count>0)
                        {
                            foreach(int sharedTermID in sharedTermIDs.Keys)
                            {
                                double weight = 0;
                                if (termFreqs1.ContainsKey(sharedTermID))
                                {
                                    weight = termFreqs1[sharedTermID].Count*termFreqs1[sharedTermID].Weight;
                                }
                                if (termFreqs2.ContainsKey(sharedTermID))
                                {
                                    weight += termFreqs2[sharedTermID].Count*termFreqs2[sharedTermID].Weight;
                                }
                                if(supportiveTermWeights.ContainsKey(sharedTermID))
                                {
                                    supportiveTermWeights[sharedTermID] = supportiveTermWeights[sharedTermID] + weight;
                                }
                                else
                                {
                                    supportiveTermWeights.Add(sharedTermID,weight);
                                }
                            }
                        }
                    }
                }
                // if single doc, the distance is 0
                if(docDistances.Count==0)
                    docDistances.Add(0);
                if(this._DocDistances.ContainsKey(docID1))
                {
                    this._DocDistances[docID1] = StatsUtil.GetAverage(docDistances).Value;
                }
                else
                {
                    this._DocDistances.Add(docID1, StatsUtil.GetAverage(docDistances).Value);
                }
            }
            if(clusterDistances.Count==0)
                clusterDistances.Add(0);

            double clusterAvgDistance = StatsUtil.GetAverage(clusterDistances).Value;
            clusterNode.Distance = clusterAvgDistance;
            clusterNode.SupportiveTermWeights = supportiveTermWeights;
        }

        private int NextClusterID()
        {
            int lastClusterID = 0;
            if (this._ClusterResult.Count > 0)
            {
                foreach (int clusterID in this._ClusterResult.Keys)
                {
                    if (clusterID > lastClusterID)
                        lastClusterID = clusterID;
                }
            }
            return lastClusterID + 1;
        }

        private Dictionary<int, int> InitTermFrequency(Dictionary<int, DocTermFreq> termFreqs)
        {
            Dictionary<int,int> termFreq=new Dictionary<int, int>();
            if(termFreqs !=null && termFreqs.Count>0)
            {
                foreach(int termID in termFreqs.Keys)
                {
                    termFreq.Add(termID,termFreqs[termID].Count);
                }
            }
            return termFreq;
        }
        #endregion

        #region update
        /// <summary>
        /// return number of docs reassociated to new clusters
        /// </summary>
        /// <returns></returns>
        private int UpdateDocAssociations()
        {
            int[] docIDs=new int[this._DocTerms.Count];
            this._DocTerms.Keys.CopyTo(docIDs,0);
            int[] clusterIDs=new int[this._ClusterResult.Count];
            this._ClusterResult.Keys.CopyTo(clusterIDs,0);
            int numberDocUpdated = 0;
            for(int i=0;i<docIDs.Length;i++)
            {
                int docID = docIDs[i];
                int oldClusterID = 0;
                for (int k = 0; k < clusterIDs.Length; k++)
                {
                    ClusterNode clusterNode = this._ClusterResult[clusterIDs[k]];
                    if (clusterNode.DocIDs.Contains(docID))
                    {
                        oldClusterID = clusterIDs[k];
                        break;
                    }
                }
                for(int k=0;k<clusterIDs.Length;k++)
                {
                    ClusterNode clusterNode = this._ClusterResult[clusterIDs[k]];
                    if (clusterIDs[k]!=oldClusterID)
                    {
                        double newDistance = this.CalculateDocClusterDistance(docID, clusterNode);
                        if (newDistance < this._DocDistances[docID])
                        {
                            ClusterNode oldClusterNode = this._ClusterResult[oldClusterID];
                            oldClusterNode.DocIDs.Remove(docID);
                            clusterNode.DocIDs.Add(docID);
                            numberDocUpdated += 1;
                            return numberDocUpdated;
                        }
                    }
                }
            }
            return numberDocUpdated;
        }

        /// <summary>
        /// the average distance of doc to other docs within cluster
        /// </summary>
        /// <param name="docID"></param>
        /// <param name="clusterNode"></param>
        /// <returns></returns>
        private double CalculateDocClusterDistance(int docID, ClusterNode clusterNode)
        {
            Dictionary<int, DocTermFreq> unionTermIDs = new Dictionary<int, DocTermFreq>();
            foreach(int docIDInCluster in clusterNode.DocIDs)
            {
                Dictionary<int, DocTermFreq> termFreqs = this._DocTerms[docIDInCluster];
                foreach(int termID in termFreqs.Keys)
                {
                    if(!unionTermIDs.ContainsKey(termID))
                    {
                        unionTermIDs.Add(termID, termFreqs[termID]);
                    }
                }
            }

            Dictionary<int, DocTermFreq> termFreqs1 = this._DocTerms[docID];
            Dictionary<int, DocTermFreq> allOtherTermFreqs = new Dictionary<int, DocTermFreq>();
            foreach (int otherDocID in clusterNode.DocIDs)
            {
                if (otherDocID != docID)
                {
                    Dictionary<int, DocTermFreq> otherTermFreqs = this._DocTerms[otherDocID];
                    foreach (int termID in otherTermFreqs.Keys)
                    {
                        if (allOtherTermFreqs.ContainsKey(termID))
                        {
                            DocTermFreq otherDocTermFreq = allOtherTermFreqs[termID];
                            otherDocTermFreq.Count += otherTermFreqs[termID].Count;
                            allOtherTermFreqs[termID] = otherDocTermFreq;
                        }
                        else
                        {
                            allOtherTermFreqs.Add(termID, otherTermFreqs[termID]);
                        }
                    }
                }
            }
            //Dictionary<int, int> termFreq1 = this.InitTermFrequency(termFreqs1);
            //Dictionary<int, int> termFreq2 = this.InitTermFrequency(unionTermIDs);
            Dictionary<int, int> sharedTermIDs = new Dictionary<int, int>();
            if (this._TermWeights == null || this._TermWeights.Count == 0)
            {
                return this._DistanceCalculator.CalculateDistance(termFreqs1, allOtherTermFreqs, ref sharedTermIDs);
            }
            else if(this._TermStemmings1==null || this._TermStemmings2==null)
            {
                return this._DistanceCalculator.CalculateDistance(
                    termFreqs1, allOtherTermFreqs, this._TermWeights,
                    true, this._NGramLen, ref sharedTermIDs);
            }
            else
            {
                return this._DistanceCalculator.CalculateDistance(
                    termFreqs1, allOtherTermFreqs, this._TermWeights,
                    this._TermStemmings1, this._TermStemmings2, ref sharedTermIDs);
            }
        }
        #endregion

        #region reassociate
        private int ReassociateDocs()
        {
            int[] docIDs = new int[this._DocTerms.Count];
            this._DocTerms.Keys.CopyTo(docIDs, 0);
            int[] clusterIDs = new int[this._ClusterResult.Count];
            this._ClusterResult.Keys.CopyTo(clusterIDs, 0);
            int numDocsReassociated = 0;
            for (int i = 0; i < docIDs.Length; i++)
            {
                int docID = docIDs[i];
                List<int> oldClusterIDs = new List<int>();
                for (int k = 0; k < clusterIDs.Length; k++)
                {
                    ClusterNode clusterNode = this._ClusterResult[clusterIDs[k]];
                    if (clusterNode.DocIDs.Contains(docID))
                    {
                        oldClusterIDs.Add(clusterIDs[k]);
                    }
                }
                for (int k = 0; k < clusterIDs.Length; k++)
                {
                    if (!oldClusterIDs.Contains(clusterIDs[k]))
                    {
                        ClusterNode clusterNode = this._ClusterResult[clusterIDs[k]];
                        double newDistance = this.CalculateDocClusterDistance(docID, clusterNode);
                        if (newDistance <= this._DocDistances[docID] * this._ReassociationDistanceRatio && 
                            newDistance <= clusterNode.Distance * this._ReassociationDistanceRatio)
                        {
                            if (!clusterNode.DocIDs.Contains(docID))
                            {
                                clusterNode.DocIDs.Add(docID);
                                numDocsReassociated += 1;
                                break;
                            }
                        }
                    }
                }
            }
            return numDocsReassociated;
        }
        #endregion
    }
}
